package com.credithc.eagleeye.plugins.gateway;
import com.credithc.eagleeye.agent.factory.EagleEyeAgentFactory;
import com.credithc.eagleeye.agent.util.HostUtil;
import com.credithc.eagleeye.constants.EagleEyeConstant;
import com.credithc.eagleeye.constants.RpcTypeEnum;
import com.credithc.eagleeye.plugins.util.LogUtil;
import com.credithc.eagleeye.thread.EagleEyeThreadContext;
import com.credithc.eagleeye.thread.TraceInfo;
import com.credithc.eagleeye.util.FastJsonConvert;
import com.credithc.eagleeye.util.MixUtil;
import com.credithc.eagleeye.vo.RpcTraceInfoVo;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.slf4j.MDC;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.Date;
import java.util.List;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR;
/**
 * 描 述: <br/>
 * 作 者: zhangjin<br/>
 * 创 建：2019年11月08日<br/>
 * 版 本：v1.0.0<br>
 * <p>
 * 历 史: (版本) 作者 时间 注释 <br/>
 * 参考ModifyRequestBodyGatewayFilterFactory，ModifyResponseBodyGatewayFilterFactory
 */
@SuppressWarnings("DuplicatedCode")
@Slf4j
public class EagleEyeGatewayFilter implements GlobalFilter, Ordered {
    private final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();
    private static int MAX_REQUEST_BODY_LENGTH = 4000;

    public EagleEyeGatewayFilter() {
        log.info("EagleEyeGatewayFilter init..");
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        return processRequest(exchange, chain);
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }

    private Mono<Void> processRequest(ServerWebExchange exchange, GatewayFilterChain chain) {
        TraceInfo traceInfo = initEagleEyeTraceInfo(exchange);
        RpcTraceInfoVo rpcTraceInfoVo = initRpcTraceInfoVo(traceInfo, exchange);
        MediaType contentType = exchange.getRequest().getHeaders().getContentType();
        //字符串的报文获取body
        if(contentType!=null && (contentType.includes(MediaType.APPLICATION_JSON)
                 ||contentType.includes(MediaType.APPLICATION_JSON_UTF8)
                 ||contentType.includes(MediaType.TEXT_PLAIN)
                 ||contentType.includes(MediaType.TEXT_XML))){
            ServerRequest serverRequest = ServerRequest.create(exchange, messageReaders);
            Mono<String> modifiedBody = serverRequest.bodyToMono(String.class)
                    .map(originRequestBody ->  {
                        //可在此处改写原始的请求body，比如进行解密等等，这里仅仅用于鹰眼记录原始报文，不做任何修改
                        StringBuilder requestBodySb = new StringBuilder();
                        requestBodySb.append(rpcTraceInfoVo.getRequestJson());
                        requestBodySb.append("requestBody:\n");
                        if (StringUtils.length(originRequestBody) > MAX_REQUEST_BODY_LENGTH) {
                            requestBodySb.append(StringUtils.left(originRequestBody,MAX_REQUEST_BODY_LENGTH)+"...");
                        }else{
                            requestBodySb.append(originRequestBody);
                        }
                        rpcTraceInfoVo.setRequestJson(requestBodySb.toString());
                        return originRequestBody ;
                    });

            BodyInserter<Mono<String>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
            HttpHeaders tempHeaders = new HttpHeaders();
            tempHeaders.putAll(exchange.getRequest().getHeaders());
            // the new content type will be computed by bodyInserter and then set in the request decorator
            tempHeaders.remove(HttpHeaders.CONTENT_LENGTH);
            CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, tempHeaders);

            return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
                ServerHttpRequestDecorator requestDecorator = decorateRequest(exchange, traceInfo, tempHeaders, outputMessage);
                ServerHttpResponse responseDecorator = decorateResponse(exchange,rpcTraceInfoVo);
                return chain.filter(exchange.mutate().request(requestDecorator).response(responseDecorator).build()).then(Mono.fromRunnable(()-> {
                    send2EagleEye(rpcTraceInfoVo,exchange.getResponse());
                }));
            }));
        }else{
            //非字符串的报文就不获取body了
            ServerHttpRequestDecorator requestDecorator = decorateRequest(exchange, traceInfo, null, null);
            ServerHttpResponse responseDecorator = decorateResponse(exchange,rpcTraceInfoVo);
            return chain.filter(exchange.mutate().request(requestDecorator).response(responseDecorator).build()).then(Mono.fromRunnable(()-> {
                send2EagleEye(rpcTraceInfoVo,exchange.getResponse());
            }));
        }
    }

    ServerHttpRequestDecorator decorateRequest(ServerWebExchange exchange, TraceInfo traceInfo, HttpHeaders tempHeaders, CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                if (tempHeaders != null) {
                    long contentLength = tempHeaders.getContentLength();
                    if (contentLength > 0) {
                        httpHeaders.setContentLength(contentLength);
                    } else {
                        httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                    }
                }
                //鹰眼header传播
                httpHeaders.set(EagleEyeConstant.HC_TRACE_ID, traceInfo.getTraceId());
                httpHeaders.set(EagleEyeConstant.HC_RPC_ID, traceInfo.getRpcId());
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                if (outputMessage != null) {
                    return outputMessage.getBody();
                }
                return super.getBody();
            }
        };
    }
    //TODO:decorateResponse
    ServerHttpResponse decorateResponse(ServerWebExchange exchange,RpcTraceInfoVo rpcTraceInfoVo) {
        return new ServerHttpResponseDecorator(exchange.getResponse()) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                Class inClass = String.class;
                Class outClass = String.class;
                String originalResponseContentType = exchange.getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);

                if(StringUtils.containsIgnoreCase(originalResponseContentType,"json")
                        || StringUtils.containsIgnoreCase(originalResponseContentType,"text")
                        || StringUtils.containsIgnoreCase(originalResponseContentType,"xml")) {
                    HttpHeaders httpHeaders = new HttpHeaders();
                    // explicitly add it in this way instead of
                    // 'httpHeaders.setContentType(originalResponseContentType)'
                    // this will prevent exception in case of using non-standard media
                    // types like "Content-Type: image"
                    httpHeaders.add(HttpHeaders.CONTENT_TYPE, originalResponseContentType);
                    ClientResponse clientResponse = ClientResponse
                            .create(exchange.getResponse().getStatusCode())
                            .headers(headers -> headers.putAll(httpHeaders))
                            .body(Flux.from(body)).build();
                    // TODO: flux or mono
                    Mono<String> modifiedBody = clientResponse.bodyToMono(String.class)
                            .map(originalBody -> {
                                //可在此处改写原始的response，比如进行加密等等
                                StringBuilder sb = new StringBuilder();
                                sb.append("responseBody:\n");
                                if (StringUtils.length(originalBody) > MAX_REQUEST_BODY_LENGTH) {
                                    sb.append(StringUtils.left(originalBody,MAX_REQUEST_BODY_LENGTH)+"...");
                                }else{
                                    sb.append(originalBody);
                                }
                                rpcTraceInfoVo.setResponseJson(sb.toString());
                                return originalBody;
                            });
                    BodyInserter<Mono<String>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
                    CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, exchange.getResponse().getHeaders());
                    return bodyInserter.insert(outputMessage, new BodyInserterContext())
                            .then(Mono.defer(() -> {
                                Flux<DataBuffer> messageBody = outputMessage.getBody();
                                HttpHeaders headers = getDelegate().getHeaders();
                                if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
                                    messageBody = messageBody.doOnNext(data -> {
                                        int contentLength = data.readableByteCount();
                                        headers.setContentLength(contentLength);
                                        rpcTraceInfoVo.setResponseSize(contentLength);
                                    });
                                }
                                // TODO: fail if isStreamingMediaType?
                                return getDelegate().writeWith(messageBody);
                            }));
                }
                return super.writeWith(body);
            }

            @Override
            public Mono<Void> writeAndFlushWith(
                    Publisher<? extends Publisher<? extends DataBuffer>> body) {
                return writeWith(Flux.from(body).flatMapSequential(p -> p));
            }
        };
    }

    private TraceInfo initEagleEyeTraceInfo(ServerWebExchange exchange) {
        ServerHttpRequest request = exchange.getRequest();
        EagleEyeThreadContext.init();
        TraceInfo traceInfo = new TraceInfo();
        String hcTraceId = request.getHeaders().getFirst(EagleEyeConstant.HC_TRACE_ID);
        if (StringUtils.isNotBlank(hcTraceId)) {
            traceInfo.setTraceId(hcTraceId);
        }
        String hcRpcId = request.getHeaders().getFirst(EagleEyeConstant.HC_RPC_ID);
        if (hcRpcId != null && hcRpcId.length() > 0) {
            traceInfo.setRpcId(hcRpcId);
            traceInfo.addHierarchy();
            traceInfo.addSequenceNo();
        }
        EagleEyeThreadContext.putTraceInfo(traceInfo);
        MDC.put("traceId", traceInfo.getTraceId());
        MDC.put("rpcId", traceInfo.getRpcId());
        MDC.put("clientIp", getClientIp(request));
        return traceInfo;
    }

    private RpcTraceInfoVo initRpcTraceInfoVo(TraceInfo traceInfo, ServerWebExchange exchange) {
        ServerHttpRequest request = exchange.getRequest();
        RpcTraceInfoVo rpcTraceInfoVo = new RpcTraceInfoVo();
        Date date = new Date();
        rpcTraceInfoVo.setRequestDateTime(MixUtil.ISO_DATETIME_TIME_ZONE_FORMAT_WITH_MILLIS.format(date));
        rpcTraceInfoVo.setTraceId(traceInfo.getTraceId());
        rpcTraceInfoVo.setRpcId(traceInfo.getRpcId());
        rpcTraceInfoVo.setRpcType(RpcTypeEnum.HTTP.getRpcName());
        rpcTraceInfoVo.setServiceCategory("SpringGateway");
        rpcTraceInfoVo.setServiceName("SpringGateway");
        String path = StringUtils.removeStart(request.getPath().pathWithinApplication().value(), "/");
        rpcTraceInfoVo.setMethodName(path);
        rpcTraceInfoVo.setClientHost(HostUtil.getIp());
        rpcTraceInfoVo.setServerHost(request.getURI().toString());
        rpcTraceInfoVo.setRequestSize(0);
        rpcTraceInfoVo.setRunTime(System.currentTimeMillis());
        // 获取路由的目标URI
        URI targetUri = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        HttpMethod method = request.getMethod();
        StringBuilder sb = new StringBuilder();
        sb.append(String.format("targetUri: %s\n", targetUri == null ? path : targetUri.toString()));
        sb.append(String.format("method: %s\n", method == null ? "" : method.name()));
        /*
        sb.append(String.format("headers size: %s\n", request.getHeaders().size()));
        sb.append("-------------\n");
        request.getHeaders().forEach((key, value) -> {
            sb.append(String.format("%s: %s\n", key, String.join(",", value)));
        });
        */
        MultiValueMap<String, String> queryParams = request.getQueryParams();
        sb.append(String.format("queryParams: %s\n", FastJsonConvert.convertObjectToJSON(queryParams)));
        rpcTraceInfoVo.setRequestJson(sb.toString());
        rpcTraceInfoVo.setRequestSize(request.getHeaders().getContentLength());
        return rpcTraceInfoVo;
    }

    private void send2EagleEye(RpcTraceInfoVo rpcTraceInfoVo, ServerHttpResponse response) {
        try {
            rpcTraceInfoVo.setRunTime(System.currentTimeMillis() - rpcTraceInfoVo.getRunTime());
            HttpStatus statusCode = response.getStatusCode();
            StringBuilder sb = new StringBuilder();
            if(statusCode!=null) {
                sb.append("statusCode:" + statusCode.toString()+"\n");
            }
            if (rpcTraceInfoVo.getResponseJson() != null) {
                sb.append(rpcTraceInfoVo.getResponseJson());
            }
            rpcTraceInfoVo.setResponseJson(sb.toString());
            if (statusCode.isError()) {
                rpcTraceInfoVo.setResult("ERROR");
            } else {
                rpcTraceInfoVo.setResult("OK");
            }
            EagleEyeAgentFactory.getInstance().send2EagleEye(rpcTraceInfoVo);
            LogUtil.printSplunkLog(rpcTraceInfoVo);
        } catch (Exception ex) {
        }
    }

    private String getClientIp(ServerHttpRequest request) {
        if (request == null) {
            return "";
        }
        String ip = request.getHeaders().getFirst("x-forwarded-for");
        if (ip == null || ip.length() == 0) {
            ip = request.getHeaders().getFirst("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0) {
            ip = request.getHeaders().getFirst("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0) {
            ip = request.getHeaders().getFirst("X-Real-IP");
        }
        if (ip == null || ip.length() == 0) {
            ip = request.getRemoteAddress().getAddress().getHostAddress();
        }
        if (ip != null) {
            int index = ip.indexOf(",");
            ip = (index > -1) ? ip.substring(0, index) : ip;
        }
        return "0:0:0:0:0:0:0:1".equals(ip) ? "127.0.0.1" : ip;
    }
}
